package db

import (
	"fmt"
	"log"
	"os"
	"strings"

	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
)

var DB *gorm.DB

func InitDB() error {
	// 获取数据库类型，默认为 mysql
	dbType := strings.ToLower(os.Getenv("DB_TYPE"))
	if dbType == "" {
		dbType = "mysql" // 默认使用 MySQL
	}

	var err error

	switch dbType {
	case "mysql":
		// MySQL 连接
		// 适配新的环境变量名称
		username := getEnvOrDefault("DB_USER", "DB_USERNAME", "root")
		password := os.Getenv("DB_PASSWORD")
		host := os.Getenv("DB_HOST")
		port := os.Getenv("DB_PORT")
		dbname := getEnvOrDefault("DB_NAME", "DB_DATABASE", "dynamic_db")

		dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
			username,
			password,
			host,
			port,
			dbname)

		DB, err = gorm.Open(mysql.Open(dsn), &gorm.Config{})
		if err != nil {
			return fmt.Errorf("failed to connect to MySQL: %v", err)
		}
		log.Printf("Connected to MySQL database: %s@%s:%s/%s", username, host, port, dbname)

	case "postgres":
		// PostgreSQL 连接
		username := getEnvOrDefault("DB_USER", "DB_USERNAME", "postgres")
		password := os.Getenv("DB_PASSWORD")
		host := os.Getenv("DB_HOST")
		port := os.Getenv("DB_PORT")
		dbname := getEnvOrDefault("DB_NAME", "DB_DATABASE", "postgres")

		dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%s sslmode=disable",
			host,
			username,
			password,
			dbname,
			port)

		DB, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
		if err != nil {
			return fmt.Errorf("failed to connect to PostgreSQL: %v", err)
		}
		log.Printf("Connected to PostgreSQL database: %s@%s:%s/%s", username, host, port, dbname)

	default:
		return fmt.Errorf("unsupported database type: %s", dbType)
	}

	log.Println("Database connection established successfully")
	return nil
}

// GetDatabaseType 返回当前连接的数据库类型
func GetDatabaseType() string {
	dbType := strings.ToLower(os.Getenv("DB_TYPE"))
	if dbType == "" {
		return "mysql" // 默认
	}
	return dbType
}

// getEnvOrDefault 获取环境变量，如果第一个不存在则尝试第二个，都不存在则返回默认值
func getEnvOrDefault(key1, key2, defaultValue string) string {
	value := os.Getenv(key1)
	if value == "" {
		value = os.Getenv(key2)
		if value == "" {
			return defaultValue
		}
	}
	return value
}
